module type Problem_Type =
sig
  val n: int
  val m: int
  val population: int
  val init: unit -> int array array
  val eval: int array -> int
  val w: float
  val c1: float
  val c2: float
end

module Make (P: Problem_Type) = 
struct
  open P
    
  (* initialisation of swarm particles and speed *)
  let x = init ()
  let v = Array.init population (fun _ -> Array.init n (fun _ -> 0))

  let best_eval = ref max_int
  let local_best_eval = Array.init population (fun _ -> max_int)

  (* best particle of the swarm *)
  let best = 
    let best = ref x.(0) in
    let e = ref (eval !best) in
      local_best_eval.(0) <- !e;
      for i=1 to population-1 do
	let v = eval x.(i) in
	  if v < !e then begin
	    best := x.(i);
	    e := v;
	  end;
	  local_best_eval.(i) <- v;
      done;
      best_eval := !e;
      Array.copy !best
	
  (*  best known position of each particle *)
  let local_best = Array.map Array.copy x

  let update_speed = fun () ->
    for i=0 to population-1 do
      for j=0 to n-1 do
	let r1 = Random.float 1. 
	and r2 = Random.float 1. in
	  v.(i).(j) <-
	    int_of_float (w *. (float v.(i).(j)) +. c1 *. r1 *. (float (best.(j) - x.(i).(j))) +. c2 *. r2 *. (float (local_best.(i).(j) - x.(i).(j))));
	  if v.(i).(j) > 0 
	  then v.(i).(j) <- min v.(i).(j) m
	  else v.(i).(j) <- max v.(i).(j) (-m)
      done
    done

  let update_position = fun () ->
    for i=0 to population-1 do
      for j=0 to n-1 do
	x.(i).(j) <- (x.(i).(j) + v.(i).(j)) mod m
      done;
      let e = eval x.(i) in
	if e < local_best_eval.(i) 
	then begin
	  local_best_eval.(i) <- e;
	  local_best.(i) <- Array.copy x.(i)
	end;
	if e < !best_eval
	then begin
	  best_eval := e;
	  for k=0 to n-1 do
	    best.(k) <- x.(i).(k);
	  done
	end
    done


  let solve = fun () ->
    let generation = ref 0 in
    while true do
      incr generation;
      update_speed ();
      update_position ();
      Printf.printf "% 4d: %d\n%!" !generation !best_eval;
    done
end
